# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment1/question5/main_calculation.py
import cmath

import matplotlib.pyplot
import torch

from grating import Grating


def _calculate_analytically(grating, order):
    period = grating.period()
    width = grating.ridge_width()
    limit3 = period / 2
    limit1 = -limit3
    limit2 = limit1 + width
    p_factor = 1j * limit3 / cmath.pi
    epsilon_r = grating.ridge_dielectric_coefficient()
    epsilon_g = grating.groove_dielectric_coefficient()

    def primitive_function(p_i):
        factor = p_factor / p_i
        return lambda p_x: factor * cmath.exp(p_x / factor)

    result = {0: epsilon_r * width + epsilon_g * (1 - width)}
    for i in range(1, order + 1):
        p = primitive_function(i)
        a = epsilon_r * (p(limit2) - p(limit1))
        a += epsilon_g * (p(limit3) - p(limit2))
        result[i] = a / period
        p = primitive_function(-i)
        a = epsilon_r * (p(limit2) - p(limit1))
        a += epsilon_g * (p(limit3) - p(limit2))
        result[-i] = a / period
    return result


def _calculate_with_fft(grating, order):
    # TODO: 确认正确性
    # 和 _calculate_analytically 的结果有明显出入
    # 因此暂时也直接用了 torch.fft ，以免自己实现 fft 出现问题更难排查
    # 当然，既然用了 PyTorch ，这样去生成 y_tensor 和创建字典就不太好
    # 不过既然有更大的问题没解决，还是暂时不去处理

    period = grating.period()
    # 需要取 0 到 p ，而不能 -p/2 到 p/2 ，两者之间相位有差别
    limit3 = period
    limit1 = 0
    x_tensor = torch.linspace(limit1, limit3, order)
    y_tensor = torch.tensor([grating.dielectric_coefficient(float(x)) for x in x_tensor])
    result = torch.fft.fft(y_tensor)
    d = {}
    for i, f in enumerate(result):
        d[i] = complex(f) / order
    return d


def _draw_from_dictionary_of_a(a: dict[int, complex], p=1, **plot_args):
    def f(x):
        items = (a[i] * cmath.exp(1j * 2 * cmath.pi * i * x / p) for i in a)
        return sum(items).real

    x_list = torch.linspace(-0.6 * p, 0.6 * p, 2000)
    y_list = [f(float(v)) for v in x_list]
    return matplotlib.pyplot.plot(x_list, y_list, **plot_args)[0]


def _draw_grating(grating: Grating, **plot_args):
    x_list = torch.linspace(-0.6 * grating.period(), 0.6 * grating.period(), 2000)
    y_list = [grating.dielectric_coefficient(float(v)) for v in x_list]
    return matplotlib.pyplot.plot(x_list, y_list, **plot_args)[0]


def _main():
    grating = Grating(0.4)
    a_a = _calculate_analytically(grating, 40)
    a_f = _calculate_with_fft(grating, 40)

    matplotlib.pyplot.rcParams['font.sans-serif'] = ['SimHei']
    matplotlib.pyplot.rcParams['axes.unicode_minus'] = False

    matplotlib.pyplot.subplot(121)
    matplotlib.pyplot.axvspan(-grating.period() / 2, grating.period() / 2, alpha=0.1)
    a_a_path = _draw_from_dictionary_of_a(a_a, grating.period())
    g_path = _draw_grating(grating)
    matplotlib.pyplot.legend(
        [g_path, a_a_path],
        ["原函数", "通过解析方式"]
        , loc="lower right")

    matplotlib.pyplot.subplot(122)
    matplotlib.pyplot.axvspan(-grating.period() / 2, grating.period() / 2, alpha=0.1)
    a_a_path = _draw_from_dictionary_of_a(a_a, grating.period())
    a_f_path = _draw_from_dictionary_of_a(a_f, grating.period())
    matplotlib.pyplot.legend(
        [a_a_path, a_f_path],
        ["通过解析方式", "通过快速傅里叶变换"]
        , loc="lower right")

    matplotlib.pyplot.show()


if __name__ == "__main__":
    _main()
